package com.ljb.jsoup.ziroom.ocr;

import com.alibaba.fastjson.JSONObject;
import com.sun.org.apache.xml.internal.security.utils.Base64;
import okhttp3.FormBody;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import org.apache.commons.lang3.StringUtils;
import org.jsoup.Connection;
import org.jsoup.Jsoup;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.List;
import java.util.Map;

/**
 * 百度ocr工具类
 */
public class BaiDuOCR {
    private static String cookies = null;
    private static final String CONTENT_TYPE = "application/x-www-form-urlencoded";
    private static final String URL_HOST = "http://ai.baidu.com/aidemo";
    private static final String URL_COOKIE = "http://ai.baidu.com/tech/ocr/general";
    //高精度版
    private static final String TYPE_ACCURATE = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate_basic";

    private BaiDuOCR() {
    }

    private static volatile BaiDuOCR instance;

    public static BaiDuOCR getDefault() {
        if (instance == null) {
            synchronized (BaiDuOCR.class) {
                if (instance == null) {
                    instance = new BaiDuOCR();
                }
            }
        }
        return instance;
    }

    /**
     * 将文件转换为Base64编码
     *
     * @param filePath
     * @return
     */
    public String encryptToBase64(String filePath) {
        try {
            byte[] b = Files.readAllBytes(Paths.get(filePath));
            return Base64.encode(b);
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 高精度识别
     *
     * @param imagePath 图片路径
     * @return 识别结果
     */
    public List<WordsResult.WordsResultBean> ocrImage(String imagePath) throws IOException {
        File file = new File(imagePath);
        if (!file.exists()) {
            System.out.println("文件不存在: " + imagePath);
            return null;
        }
        if (!StringUtils.endsWithAny(imagePath, "jpg", "png", "jpeg", "bmp", "JPG", "PNG", "JPEG", "BMP")) {
            System.out.println("仅支持jpg,jpeg,png,bmp格式的图片");
            return null;
        }
        int index = imagePath.lastIndexOf('.');
        String type = imagePath.substring(index + 1);
        String base64 = "data:image/%s;base64,";
        base64 = String.format(base64, type.toLowerCase());
        base64 += encryptToBase64(imagePath);
        return ocr(null, base64);
    }

    /**
     * 高精度识别
     *
     * @param imageUrl 图片地址
     * @return 识别结果
     * @throws IOException
     */
    public List<WordsResult.WordsResultBean> ocrImageUrl(String imageUrl) throws IOException {
        return ocr(imageUrl, null);
    }


    private synchronized List<WordsResult.WordsResultBean> ocr(String imageUrl, String image) throws IOException {
        if (cookies == null) {
            cookies = getCookies();
        }

        FormBody.Builder bodyBuilder = new FormBody.Builder()
                .add("type", TYPE_ACCURATE);

        //判断图片资源来源,优先使用url地址
        if (StringUtils.isNotBlank(imageUrl)) {
            bodyBuilder.add("image_url", imageUrl);
        } else if (StringUtils.isNotBlank(image)) {
            bodyBuilder.add("image", image);
        } else {
            return null;
        }

        Request request = new Request.Builder()
                .addHeader("Cookie", cookies)
                .addHeader("Content-Type", CONTENT_TYPE)
                .url(URL_HOST)
                .post(bodyBuilder.build())
                .build();

        OkHttpClient okHttpClient = new OkHttpClient.Builder()
                .build();

        long start = System.currentTimeMillis();
        Response response = okHttpClient
                .newCall(request)
                .execute();
        String body = response.body().string();
        long end = System.currentTimeMillis();
//        System.out.println("OCR图片识别,用时:" + (end - start) + "ms ,ret：" + body);

        JSONObject ret = JSONObject.parseObject(body);
        String errno = ret.get("errno").toString();

        //识别成功
        if (StringUtils.equals("0", errno)) {
            String data = ret.get("data").toString();
            return JSONObject.parseObject(data, WordsResult.class).words_result;

            //cookies过期
        } else if (StringUtils.equals("102", errno)) {
            //重新获取cookies
            cookies = null;
            return ocr(imageUrl, image);
        }
        return null;
    }

    /**
     * 获取百度AI网页的cookie信息
     */
    private String getCookies() throws IOException {
        Connection.Response response = Jsoup.connect(URL_COOKIE).execute();
        Map<String, String> cookies = response.cookies();
        String cookies_string = cookies.toString();
        cookies_string = cookies_string.substring(cookies_string.indexOf("{") + 1, cookies_string.lastIndexOf("}"));
        cookies_string = cookies_string.replaceAll(",", ";");
        System.out.println("获取到的Cookie为:" + cookies_string);
        return cookies_string;
    }
}
